# -*- coding: utf-8 -*-

from __future__ import print_function
import sys
sys.path.append('/opt/alps/lib')
sys.path.append('/share/opt/alps/lib')
import pyalps
import matplotlib.pyplot as plt
import pyalps.plot
import numpy as np
import time
#import scipy.io
import glob, os, shutil
import multiprocessing
from copy import deepcopy

#----------------------1 Run the Task------------------------

run_with_GUI = 1
run_the_task = 1
numcores = 4
#Paras = [32,64,96,128]
Paras = [50, 100, 150]

time1 = time.time()
task_name = 'mps_spin_half'  # global variable, can be accessed by function Runmps()
sweeps = 10
L = 32

#maxstates = 300


def Runmps(para):
    maxstates = para  ##infomation hidden,ie (maxstates,sweeps) = para
    parm = {}
    parm['LATTICE'] = 'open chain lattice'
    parm['L'] = L
    parm['MODEL'] = 'spin'
    parm['CONSERVED_QUANTUMNUMBERS'] = 'Sz'
    parm['Sz_total'] = 0
    parm['J'] = 1
    parm['MAXSTATES'] = maxstates
    parm['SWEEPS'] = sweeps
    parm['NUMBER_EIGENVALUES'] = 2
    #    parm['MEASURE_LOCAL[Sz]'] = 'Sz'
    parm = [parm]
    input_file = pyalps.writeInputFiles(task_name + '_para=' + str(para), parm)
    pyalps.runApplication('mps_optim', input_file, writexml=True)

    if run_with_GUI == 0:
        doc1 = open('1Run_print.txt', 'a')
        print('para = %s task finished.' % (para), file=doc1)
        doc1.close()


def delete_old_outfiles():  # 删除上一次运行得到的文件
    old_in_files = glob.glob('*%s*.in.*' % task_name)
    old_out_files = glob.glob('*%s*.out.*' % task_name)
    for i in old_in_files + old_out_files:
        try:
            os.remove(i)  #不能删文件夹
        except OSError:
            try:
                shutil.rmtree(i)  #用来删文件夹
            except:
                print('Something wrong when deleting files.')


##如果不需要接着计算，删除占用大的MPS波函数chkp文件夹,就运行这个函数
def delete_MPS_wavefunction():
    chkp_files = glob.glob('*%s*.out.*.chkp' % task_name)
    for i in chkp_files:
        try:
            shutil.rmtree(i)  #用来删文件夹
        except:
            print('Maybe no such file to delete.')
    print('check point files deleted.')


if run_the_task == 1:
    delete_old_outfiles()  # 决定是否删除上一次运行得到的文件

    pool = multiprocessing.Pool(processes=numcores)
    ##    Paras = [i+1 for i in range(20)]
    print('report: This task has %d parallel task/tasks.' % (len(Paras)))
    print('task %s is Running...' % (task_name))
    pool.map(Runmps, Paras)
    pool.close()
    pool.join()
    time2 = time.time()
    print('Finished.Running mps has spent %.1f seconds.' % (time2 - time1))
    if run_with_GUI == 0:
        doc1 = open('1Run_print.txt', 'a')
        print(
            'Finished. Running mps has spent %.1f seconds.' % (time2 - time1),
            file=doc1)
        doc1.close()
    delete_MPS_wavefunction()  # 决定是否删除chkp文件夹
else:
    print('Task did not run, continue with data analysis.')

#----------------------2 结果分析----------------------------------------------
## 用spyder分析时请注释掉上面写runApplication的部分，这是我想到最好的办法了

result_files = pyalps.getResultFiles(prefix=task_name)
assert result_files != [], 'Report: Error, no result file has been calculated out. The calculation has definitly failed.'
## 如果没有产生结果文件，马上报错，而不是隐藏着继续运行
if run_with_GUI == 0:
    plt.switch_backend('pdf')
## 解决无GUI界面不能正常画图的问题

## 适用于both多能级和单能级的get obserable,适合local observable，Energy和Truncated_Weight
def get_get_obervable_v3(propstring, observablestr):  ## for any energy levels
    import numpy
    def load_truncated_weight_for_dset_general_level(DataSet):
        ## 得到truncated_weight的核心:
        def load_truncated_weight_for_dset(filename):
            #根据filename在h5文件中强行找到这个DataSet对应的truncated_weight
            ar = pyalps.hdf5.archive(filename)
            sweeps = ar.list_children('/spectrum/iteration')
            sweeps = [int(s) for s in sweeps]
            max_sweep = max(sweeps)
            truncated_weight_list = ar['/spectrum/iteration/' + str(max_sweep) +
                                       '/results/TruncatedWeight/mean/value']
            truncated_weight = max(truncated_weight_list)
            return truncated_weight

        num_levels = len(DataSet.y)
        truncated_weights = []
        if num_levels>1:
            for i in range(num_levels):
                filename = DataSet.props['filename']
                ff = filename.split('.')
                ff.insert(-1, str(i))
                filename = '.'.join(ff)
                truncated_weight = load_truncated_weight_for_dset(filename)
                truncated_weights.append(truncated_weight)
        else:
            filename = DataSet.props['filename']
            truncated_weight = load_truncated_weight_for_dset(filename)
            truncated_weights.append(truncated_weight)
        truncated_weights = numpy.array(truncated_weights)
        return truncated_weights

    if observablestr == 'Truncated_Weight':
        eigen_measure_obs = pyalps.loadEigenstateMeasurements(
            result_files, what=['Energy'])
        for i in eigen_measure_obs:
            DataSet = i[0]
            truncated_weights = load_truncated_weight_for_dset_general_level(
                DataSet)
            i[0].y = truncated_weights
            i[0].props['observable'] = 'Truncated_Weight'
    else:
        eigen_measure_obs = pyalps.loadEigenstateMeasurements(
            result_files, what=[observablestr])
    ## 需要各组参数计算，都拥有相同数目的能级,依赖alps从低能级到高能级排
    observables = []
    num_levels = len(eigen_measure_obs[0][0].y)
    for i in range(num_levels):
        eigen_measure_obs_i = deepcopy(eigen_measure_obs)
        for j in eigen_measure_obs_i:
            j[0].y = [j[0].y[i]]
        paras_vs_observables_i = pyalps.collectXY(
            eigen_measure_obs_i, x=propstring, y=observablestr)
        paras = paras_vs_observables_i[0].x
        observables_i = paras_vs_observables_i[0].y
        observables.append(observables_i)
    if num_levels == 1:
        observables = observables[0]
    observables = numpy.array(observables)
    return (paras, observables)


#(maxstates, energy) = get_get_obervable_v3('MAXSTATES', 'Energy')
#print(maxstates)
#print(energy)

#(maxstates, truncation_error) = get_get_obervable_v3('MAXSTATES',
#                                                     'Truncated_Weight')
#print(maxstates)
#print(truncation_error)


## get energy gap
def get_gap(propstring):
    (paras,energy) = get_get_obervable_v3(propstring, 'Energy')
    assert len(energy)==2, 'calculated energy levels are not 2'
    Egaps = energy[1]-energy[0]
    return (paras, Egaps)


#(Ls, energy_gaps) = get_gap('L')
#print(Ls)
#print(energy_gaps)
#np.savez(
#    'maxstates%s_moreLs_vs_gaps_data' % (maxstates),
#    Ls=Ls,
#    energy_gaps=energy_gaps)

(maxstates, energy_gaps) = get_gap('MAXSTATES')
print(maxstates)
print(energy_gaps)
#np.savez(
#    'L%s_maxstates_vs_gaps_data.npz' % (L),
#    maxstates=maxstates,
#    energy_gaps=energy_gaps)

## 画L_vs_gaps图
#maxstates = [50,100,150,200,250]
#maxstates = [50,300]
#
#fig = plt.figure()
#(fig, ax) = plt.subplots(figsize=(4, 3), dpi=100)
#
#for i in [50,100,150,200,300]:
#    data = np.load('maxstates%s_Ls_gaps_data.npz'%(i))
#    Ls = data['Ls']
#    energy_gaps = data['energy_gaps']
#    print(energy_gaps)
#    ax.plot(1/Ls,energy_gaps,'--',label='D=%s'%(i))
#
#ax.legend(loc=1,fontsize=8)
#ax.set_xlabel('L')
#ax.set_ylabel('Energy Gaps')
##ax.set_ylim([1e-8,1e2])
#ax.grid()
#plt.tight_layout() ##solve figure in pdf saved cutted off
#fig.savefig('L_vs_gaps.pdf')

### extrapolate energy gap to maxstates->Inf limit
#maxstates = [50,100,150,200,250]
#GapL32=[]
#GapL64=[]
#GapL96=[]
#GapL128=[]
#for i in maxstates:
#    data = np.load('maxstates%s_Ls_gaps_data.npz'%(i))
#    Ls = data['Ls']
#    energy_gaps = data['energy_gaps']
##    ax.plot(1/Ls,energy_gaps,'.-',label='D=%s'%(i))
#
#    GapL32.append(energy_gaps[0])
#    GapL64.append(energy_gaps[1])
#    GapL96.append(energy_gaps[2])
#    GapL128.append(energy_gaps[3])
#GapDiffL32 = np.diff(GapL32)
#GapDiffL64 = np.diff(GapL64)
#GapDiffL96 = np.diff(GapL96)
#GapDiffL128 = np.diff(GapL128)
##scipy.io.savemat('GapDiff_data.mat',{'GapDiffL32':GapDiffL32,'GapDiffL64':GapDiffL64,
##                                     'GapDiffL96':GapDiffL96,'GapDiffL128':GapDiffL128})
#fig = plt.figure()
#(fig, ax) = plt.subplots(figsize=(4, 3), dpi=100)
#ax.plot(np.array([1,2,3,4]),GapDiffL32)
##ax.legend(loc=1,fontsize=8)
#ax.set_xlabel('D number')
#ax.set_ylabel('Gap Difference with last')
##ax.set_yscale('log')
##ax.set_ylim([1e-8,1e2])
#ax.grid()
##ax.set_title('SPINmodel MPS convergence')

#### 得到能量
##(maxstates,Energies)=get_get_obervable_v3('MAXSTATES','Energy')
#(Ls, Energies) = get_obervable('L', 'Energy')
#print('Maxstates is %s' % (maxstates))
#print(Ls)
#print(Energies)
#e0 = (Energies[1]-Energies[0])/2
#print('e0 is %s'%(e0))

### pearson correaltion coefficient and linear fit
#pearson_coefficient_matrix = np.corrcoef(Ls,Energies)
#pearson_coefficient = pearson_coefficient_matrix[0,1]
#print('Linear correaltion coefficient of the data is %s'%(pearson_coefficient))
#polydata = np.polyfit(Ls,Energies,1)
#print('The slope of the line is %s'%(polydata[0]))

#polyfunc = np.poly1d(polydata)
#Ls2 = np.linspace(0,200,201)
#Energies2 = polyfunc(Ls2)
#fig = plt.figure()
#(fig, ax) = plt.subplots(figsize=(4, 3), dpi=300)
#ax.plot(Ls2,Energies2)
#ax.set_xlabel('Ls')
#ax.set_ylabel('Energies')
#ax.grid()
#ax.set_title('Ls Energies fitted curve')
#plt.tight_layout()  ##solve figure in pdf saved cutted off

#(sweepses,Energies) = get_obervable('SWEEPS','Energy')
#print(sweepses)
#print(Energies)
#
#fig = plt.figure()
#(fig, ax) = plt.subplots(figsize=(4, 3), dpi=300)
#ax.plot(sweepses,Energies-Energies[-1])
#ax.set_xlabel('sweeps')
#ax.set_ylabel('Energie difference')
#ax.set_yscale('log')
##ax.set_ylim([1e-8,1e2])
#ax.grid()
#ax.set_title('Energy difference convergence with sweeps')
#plt.tight_layout()  ##solve figure in pdf saved cutted off
#fig.savefig('Energy_difference_convergence_with_sweeps.pdf')










## 封存
#def get_obervable_v1(propstring, observablestr):  ##for 1 energy level
#    #从结果文件中extract观测量随参数的变化数据,依赖全局变量result_files
#    eigen_measure_obs = pyalps.loadEigenstateMeasurements(
#        result_files, what=[observablestr])
#    paras_vs_observables = pyalps.collectXY(
#        eigen_measure_obs, x=propstring, y=observablestr)
#    paras = paras_vs_observables[0].x
#    observables = paras_vs_observables[0].y
#    return (paras, observables)